import sys
import os
import time
import tensorflow as tf
import numpy as np
from modelscope.models.audio.tts import SambertHifigan
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

#model_dir = os.path.abspath("./pretrain_yangmi")
model_dir = os.path.abspath(sys.argv[1])

custom_infer_abs = {
    'voice_name':
    'F7',
    'am_ckpt':
    os.path.join(model_dir, 'tmp_am', 'ckpt'),
    'am_config':
    os.path.join(model_dir, 'tmp_am', 'config.yaml'),
    'voc_ckpt':
    os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan', 'ckpt'),
    'voc_config':
    os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan',
             'config.yaml'),
    'audio_config':
    os.path.join(model_dir, 'data', 'audio_config.yaml'),
    'se_file':
    os.path.join(model_dir, 'data', 'se', 'se.npy')
}
kwargs = {'custom_ckpt': custom_infer_abs}

model_id = SambertHifigan(os.path.join(model_dir, "orig_model"), **kwargs)

inference = pipeline(task=Tasks.text_to_speech, model=model_id, device='cpu')
srctxt = sys.argv[2]
dstpath = sys.argv[3]
with open(srctxt, 'r') as f:
   alltxt = f.readlines() 
for txtline in alltxt:
    txtline=txtline.strip().split()
    start = time.time()
    output = inference(input=txtline[1])
    print('spend time {}'.format(time.time()-start))
    wav = output[OutputKeys.OUTPUT_WAV]
    with open(os.path.join(dstpath,txtline[0]+".wav"), 'wb') as f:
        f.write(wav)

#print('am para')
#am_model_para = model_id.voices['F7'].am.state_dict()
#for k,v in am_model_para.items():
#     print("{:20s} {}".format(k, v.shape))

print('voc para')
voc_model_para = model_id.voices['F7'].voc_model.state_dict()
for k,v in voc_model_para.items():
     print("{:20s} {}".format(k, v.shape))

def binarize(mel, threshold=0.6):
    mask = mel[:,:,-1] < threshold
    vflag = tf.where(mask,tf.zeros([1]),tf.ones([1]))
    res_mel = tf.concat([mel[:,:,:-1],tf.expand_dims(vflag,axis=-1)],axis=-1)
    return res_mel

def voc_infer(post_mel):

    class UpsamModule(tf.keras.layers.Layer):
        def __init__(self,upsample_ratio,**kwargs):
            super(UpsamModule, self).__init__(**kwargs)
            self.upsample_ratio = upsample_ratio

        def call(self,x):
            return tf.repeat(x,repeats=[self.upsample_ratio],axis=1)
    class ResidualBlock(tf.keras.layers.Layer):
        def __init__(self, channels,kernel_size, dilation=[1,3,5,7],**kwargs):
            super(ResidualBlock, self).__init__(**kwargs)

            self.convs1 = []
            for li in range(4):
                cv1 = tf.keras.Sequential()
                cv1.add(tf.keras.layers.LeakyReLU(alpha=0.1))
                cv1.add(tf.keras.layers.Conv1D(channels,kernel_size,strides=1,padding='same',dilation_rate=dilation[li]))
                self.convs1.append(cv1)

            self.convs2 = []
            for li in range(4):
                cv2 = tf.keras.Sequential()
                cv2.add(tf.keras.layers.LeakyReLU(alpha=0.1))
                cv2.add(tf.keras.layers.Conv1D(channels,kernel_size,strides=1,padding='same',dilation_rate=1))
                self.convs2.append(cv2)

        def call(self,x):
            for i in range(4):
                xt = self.convs1[i](x)
                xt = self.convs2[i](xt)
                x = xt + x
            return x

    class SourceModule(tf.keras.layers.Layer):
        def __init__(self,nb_harmonics, upsample_ratio, sampling_rate, alpha=0.1, sigma=0.003,**kwargs):
            super(SourceModule, self).__init__(**kwargs)
            self.nb_harmonics = nb_harmonics
            self.upsample_ratio = upsample_ratio
            self.sampling_rate = sampling_rate
            self.alpha = alpha
            self.sigma = sigma

            self.ffn = tf.keras.layers.Conv1D(1,1,padding='same',activation=tf.keras.activations.tanh)
            #self.ups = tf.keras.layers.UpSampling1D(size=upsample_ratio)
            self.ups = UpsamModule(upsample_ratio)

        def call(self,pitch, uv):
            pitch_samples = self.ups(pitch)
            uv_samples = self.ups(uv)
            F_mat_a = []
            for i in range(self.nb_harmonics+1):
                F_mat_a.append(pitch_samples * (i + 1) / self.sampling_rate)
            F_mat = tf.concat(F_mat_a,axis=-1)
            theta_mat = 2 * np.pi * (tf.math.cumsum(F_mat, axis=1) % 1)
            u_dist = tf.random.uniform(shape=(1,1,self.nb_harmonics+1),minval=-np.pi, maxval=np.pi)
            phase_vec = tf.concat([tf.zeros([1,1,1]),u_dist[:,:,1:]],axis=-1)

            noise = tf.random.normal(shape=(1,tf.shape(pitch_samples)[1],self.nb_harmonics+1),mean=0.0,stddev=self.sigma)
            e_voice = self.alpha * tf.math.sin(theta_mat + phase_vec) + noise
            e_unvoice = self.alpha / 3 / self.sigma * noise
            e = e_voice * uv_samples + e_unvoice * (1 - uv_samples)
            return self.ffn(e)


    class Generator0(tf.keras.layers.Layer):
        def __init__(self, channels=256,kernel_size=7, upsample_kernal_sizes=[20, 11, 4, 4],upsample_scales=[10, 5, 2, 2],resblock_kernel_sizes=[3, 7, 11],**kwargs):
            super(Generator0, self).__init__(**kwargs)

            self.num_upsamples = len(upsample_kernal_sizes)
            self.num_kernels = len(resblock_kernel_sizes)

            self.conv_pre = tf.keras.layers.Conv1D(channels,kernel_size,padding='same')
            self.source_module = SourceModule(7,np.cumprod(upsample_scales)[-1],16000)

        def call(self, x):
            x = binarize(x)
            mel = x[:, :, :-2]
            pitch = x[:, :, -2:-1]
            uv = x[:, :, -1:]
            excitation = self.source_module(pitch, uv)

            x = self.conv_pre(mel)
            return x,excitation

    class Generator(tf.keras.layers.Layer):
        def __init__(self, layer,channels=256,kernel_size=7, upsample_kernal_sizes=[20, 11, 4, 4],upsample_scales=[10, 5, 2, 2],resblock_kernel_sizes=[3, 7, 11],**kwargs):
            super(Generator, self).__init__(**kwargs)

            self.num_upsamples = len(upsample_kernal_sizes)
            self.num_kernels = len(resblock_kernel_sizes)
            self.layer = layer
            self.transpose_upsamples = tf.keras.Sequential()
            self.transpose_upsamples.add(tf.keras.layers.LeakyReLU(alpha=0.1))
            self.transpose_upsamples.add(tf.keras.layers.Conv1DTranspose(channels//(2**self.layer),upsample_kernal_sizes[self.layer-1],upsample_scales[self.layer-1],padding='same'))

            self.repeat_upsamples = tf.keras.Sequential()
            self.repeat_upsamples.add(tf.keras.layers.LeakyReLU(alpha=0.1))
            self.repeat_upsamples.add(UpsamModule(upsample_scales[self.layer-1]))
            #self.repeat_upsamples.add(tf.keras.layers.LeakyReLU(alpha=0.1))
            #self.repeat_upsamples.add(tf.keras.layers.UpSampling1D(size=upsample_scales[self.layer-1]))
            self.repeat_upsamples.add(tf.keras.layers.Conv1D(channels//(2**self.layer),kernel_size,strides=1,padding='same'))

            self.conv_blocks = []
            for lj in range(self.num_kernels):
                self.conv_blocks.append(ResidualBlock(channels//(2**self.layer),resblock_kernel_sizes[lj]))


            downsample_rates = [1] + upsample_scales[::-1][:-1]
            downsample_cum_rates = np.cumprod(downsample_rates)[::-1].tolist()
            if self.layer == 4: 
                self.source_downs = tf.keras.layers.Conv1D(channels//(2**self.layer),kernel_size=1,strides=1,padding='same')
            else:
                self.source_downs = tf.keras.layers.Conv1D(channels//(2**self.layer),kernel_size=2*downsample_cum_rates[self.layer-1],strides=downsample_cum_rates[self.layer-1],padding='same')

        def call(self, x,excitation):
            x = tf.math.sin(x) + x
            rep = self.repeat_upsamples(x)
            up = self.transpose_upsamples(x)
            e = self.source_downs(excitation)
            x = rep + e + up[:, :tf.shape(rep)[1], :]

            #xs = tf.zeros([1])
            #for j in range(self.num_kernels):
            #    xs += self.conv_blocks[j](x)
            xs = self.conv_blocks[0](x)
            xs += self.conv_blocks[1](x)
            xs += self.conv_blocks[2](x)
            x = xs / self.num_kernels
            return x

    class Generator5(tf.keras.layers.Layer):
        def __init__(self, channels=256,kernel_size=7, upsample_kernal_sizes=[20, 11, 4, 4],upsample_scales=[10, 5, 2, 2],resblock_kernel_sizes=[3, 7, 11],**kwargs):
            super(Generator5, self).__init__(**kwargs)

            self.num_upsamples = len(upsample_kernal_sizes)
            self.num_kernels = len(resblock_kernel_sizes)
            self.conv_post = tf.keras.layers.Conv1D(1,kernel_size,padding='same')

            self.leac = tf.keras.layers.LeakyReLU(alpha=0.01)

        def call(self, x):
            x = self.leac(x)
            x = self.conv_post(x)
            x = tf.keras.activations.tanh(x)
            return x
    hifi0_input = tf.keras.Input(shape=(None,82),batch_size=1,name='hifi0_in')
    hifi0_out,exci = Generator0(name='hifigan0')(hifi0_input)
    hifi0_model = tf.keras.Model((hifi0_input),(hifi0_out,exci))
    hifi0_model.summary()
    hifi0_model_layer = hifi0_model.get_layer(name='hifigan0')
    timelen = 450

    #hifi1_input = tf.keras.Input(shape=(timelen,256),batch_size=1,name='hifi1_in')
    #hifi1_exci = tf.keras.Input(shape=(timelen*200,1),batch_size=1,name='hifi1_exci')
    #hifi1_input = tf.keras.Input(shape=(48,256),batch_size=1,name='hifi1_in')
    #hifi1_exci = tf.keras.Input(shape=(48*200,1),batch_size=1,name='hifi1_exci')
    hifi1_input = tf.keras.Input(shape=(21,256),batch_size=1,name='hifi1_in')
    hifi1_exci = tf.keras.Input(shape=(21*200,1),batch_size=1,name='hifi1_exci')
    hifi1_out = Generator(1,name='hifigan1')(hifi1_input,hifi1_exci)
    hifi1_model = tf.keras.Model((hifi1_input,hifi1_exci),(hifi1_out))
    hifi1_model.summary()
    hifi1_model_layer = hifi1_model.get_layer(name='hifigan1')

    #hifi2_input = tf.keras.Input(shape=(timelen*10,128),batch_size=1,name='hifi2_in')
    #hifi2_exci = tf.keras.Input(shape=(timelen*200,1),batch_size=1,name='hifi2_exci')
    #hifi2_input = tf.keras.Input(shape=(420,128),batch_size=1,name='hifi2_in')
    #hifi2_exci = tf.keras.Input(shape=(42*200,1),batch_size=1,name='hifi2_exci')
    hifi2_input = tf.keras.Input(shape=(150,128),batch_size=1,name='hifi2_in')
    hifi2_exci = tf.keras.Input(shape=(150*20,1),batch_size=1,name='hifi2_exci')
    hifi2_out = Generator(2,name='hifigan2')(hifi2_input,hifi2_exci)
    hifi2_model = tf.keras.Model((hifi2_input,hifi2_exci),(hifi2_out))
    hifi2_model.summary()
    hifi2_model_layer = hifi2_model.get_layer(name='hifigan2')

    #hifi3_input = tf.keras.Input(shape=(timelen*10*5,64),batch_size=1,name='hifi3_in')
    #hifi3_exci = tf.keras.Input(shape=(timelen*200,1),batch_size=1,name='hifi3_exci')
    #hifi3_input = tf.keras.Input(shape=(2050,64),batch_size=1,name='hifi3_in')
    #hifi3_exci = tf.keras.Input(shape=(41*200,1),batch_size=1,name='hifi3_exci')
    hifi3_input = tf.keras.Input(shape=(700,64),batch_size=1,name='hifi3_in')
    hifi3_exci = tf.keras.Input(shape=(700*4,1),batch_size=1,name='hifi3_exci')
    hifi3_out = Generator(3,name='hifigan3')(hifi3_input,hifi3_exci)
    hifi3_model = tf.keras.Model((hifi3_input,hifi3_exci),(hifi3_out))
    hifi3_model.summary()
    hifi3_model_layer = hifi3_model.get_layer(name='hifigan3')

    #hifi4_input = tf.keras.Input(shape=(timelen*10*5*2,32),batch_size=1,name='hifi4_in')
    #hifi4_exci = tf.keras.Input(shape=(timelen*200,1),batch_size=1,name='hifi4_exci')
    #hifi4_input = tf.keras.Input(shape=(4100,32),batch_size=1,name='hifi4_in')
    #hifi4_exci = tf.keras.Input(shape=(41*200,1),batch_size=1,name='hifi4_exci')
    hifi4_input = tf.keras.Input(shape=(1320,32),batch_size=1,name='hifi4_in')
    hifi4_exci = tf.keras.Input(shape=(1320*2,1),batch_size=1,name='hifi4_exci')
    hifi4_out = Generator(4,name='hifigan4')(hifi4_input,hifi4_exci)
    hifi4_model = tf.keras.Model((hifi4_input,hifi4_exci),(hifi4_out))
    hifi4_model.summary()
    hifi4_model_layer = hifi4_model.get_layer(name='hifigan4')

    hifi5_input = tf.keras.Input(shape=(None,16),batch_size=1,name='hifi5_in')
    hifi5_out = Generator5(name='hifigan5')(hifi5_input)
    hifi5_model = tf.keras.Model((hifi5_input),(hifi5_out))
    hifi5_model.summary()
    hifi5_model_layer = hifi5_model.get_layer(name='hifigan5')
    #weights = hifi_model.weights
    #for weight in weights:
    #   print('weight {}'.format(weight.name))
    #   print('weight {}'.format(weight.shape))
    #   #print('weight {}'.format(weight))
    trans_para = []
    for i in range(4):
        trans_conv_w = np.transpose(voc_model_para["transpose_upsamples.{}.1.deconv.weight".format(i)].cpu().numpy(),(2,1,0))
        trans_conv_b = voc_model_para["transpose_upsamples.{}.1.deconv.bias".format(i)].cpu().numpy()
        trans_para.extend([trans_conv_w,trans_conv_b])
    rep_para = []
    for i in range(4):
        rep_conv_w = np.transpose(voc_model_para["repeat_upsamples.{}.2.conv1d.weight".format(i)].cpu().numpy(),(2,1,0))
        rep_conv_b = voc_model_para["repeat_upsamples.{}.2.conv1d.bias".format(i)].cpu().numpy()
        rep_para.extend([rep_conv_w,rep_conv_b])
    conv_para = []
    for i in range(12):
        for j in range(4):
            conv1_conv_w = np.transpose(voc_model_para["conv_blocks.{}.convs1.{}.conv1d.weight".format(i,j)].cpu().numpy(),(2,1,0))
            conv1_conv_b = voc_model_para["conv_blocks.{}.convs1.{}.conv1d.bias".format(i,j)].cpu().numpy()
            conv_para.extend([conv1_conv_w,conv1_conv_b])
        for j in range(4):
            conv2_conv_w = np.transpose(voc_model_para["conv_blocks.{}.convs2.{}.conv1d.weight".format(i,j)].cpu().numpy(),(2,1,0))
            conv2_conv_b = voc_model_para["conv_blocks.{}.convs2.{}.conv1d.bias".format(i,j)].cpu().numpy()
            conv_para.extend([conv2_conv_w,conv2_conv_b])

    pre_conv_w = np.transpose(voc_model_para["conv_pre.conv1d.weight"].cpu().numpy(),(2,1,0))
    pre_conv_b = voc_model_para["conv_pre.conv1d.bias"].cpu().numpy()
    post_conv_w = np.transpose(voc_model_para["conv_post.conv1d.weight"].cpu().numpy(),(2,1,0))
    post_conv_b = voc_model_para["conv_post.conv1d.bias"].cpu().numpy()
    source_conv_w = np.transpose(voc_model_para["source_module.ffn.0.weight"].cpu().numpy(),(2,1,0))
    source_conv_b = voc_model_para["source_module.ffn.0.bias"].cpu().numpy()
    sdown_para = []
    for i in range(4):
        sdown_conv_w = np.transpose(voc_model_para["source_downs.{}.conv1d.weight".format(i)].cpu().numpy(),(2,1,0))
        sdown_conv_b = voc_model_para["source_downs.{}.conv1d.bias".format(i)].cpu().numpy()
        sdown_para.extend([sdown_conv_w,sdown_conv_b])

    hifi0_para = [pre_conv_w,pre_conv_b,source_conv_w,source_conv_b]
    hifi0_model_layer.set_weights(hifi0_para)

    hifi1_para = trans_para[:2] + rep_para[:2] + conv_para[:48] + sdown_para[:2]
    hifi1_model_layer.set_weights(hifi1_para)
    hifi2_para = trans_para[1*2:2*2] + rep_para[1*2:2*2] + conv_para[48:2*48] + sdown_para[1*2:2*2]
    hifi2_model_layer.set_weights(hifi2_para)
    hifi3_para = trans_para[2*2:3*2] + rep_para[2*2:3*2] + conv_para[2*48:3*48] + sdown_para[2*2:3*2]
    hifi3_model_layer.set_weights(hifi3_para)
    hifi4_para = trans_para[3*2:4*2] + rep_para[3*2:4*2] + conv_para[3*48:4*48] + sdown_para[3*2:4*2]
    hifi4_model_layer.set_weights(hifi4_para)

    hifi5_para = [post_conv_w,post_conv_b]
    hifi5_model_layer.set_weights(hifi5_para)
    #hifi_model.save('tfmodel/hifigan')

    #hifi0_model.save('tfmodel/hifigan0')
    #hifi1_model.save('tfmodel/hifigan1')
    #hifi2_model.save('tfmodel/hifigan2')
    #hifi3_model.save('tfmodel/hifigan3')
    #hifi4_model.save('tfmodel/hifigan4')
    #hifi5_model.save('tfmodel/hifigan5')

    #hifi_in = tf.expand_dims(tf.convert_to_tensor(post_mel.cpu().numpy()),axis=0)
    hifi_in = tf.expand_dims(tf.convert_to_tensor(post_mel),axis=0)
    hifi0_out,exci_out = hifi0_model((hifi_in),training=False)
    print('hifi0 {} {}'.format(hifi0_out.shape,exci_out.shape))

    #hifi1_out = hifi1_model((hifi0_out,exci),training=False)
    #print('hifi1 {}'.format(hifi1_out.shape))
    #hifi2_out = hifi2_model((hifi1_out,exci),training=False)
    #print('hifi2 {}'.format(hifi2_out.shape))
    #hifi3_out = hifi3_model((hifi2_out,exci),training=False)
    #print('hifi3 {}'.format(hifi3_out.shape))
    #hifi4_out = hifi4_model((hifi3_out,exci),training=False)
    #print('hifi4 {}'.format(hifi4_out.shape))

    #hifi5_out = hifi5_model((hifi4_out),training=False)
    #print('hifi5 {}'.format(hifi5_out.shape))

    #hifi_out_list = []
    #exci_out_list = []
    #time_len = tf.shape(hifi_in).numpy()[1]
    #for step in range(4,time_len,40):
    #    melstep = hifi_in[:,step-4:step+44,:]
    #    rpad = 48 - tf.shape(melstep).numpy()[1]
    #    if rpad > 0:
    #        melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
    #    hifi_out_step,exci_out_step = hifi0_model(melstep)
    #    print('hifi_out {} {}'.format(hifi_out_step.shape,exci_out_step.shape))
    #    if step == 4:
    #        hifi_out_step = hifi_out_step[:,:-4,:]
    #        exci_out_step = exci_out_step[:,:-4*200,:]
    #    else:
    #        if rpad == 0:
    #            hifi_out_step = hifi_out_step[:,4:-4,:]
    #            exci_out_step = exci_out_step[:,4*200:-4*200,:]
    #        else:
    #            hifi_out_step = hifi_out_step[:,4:-(rpad):,:]
    #            exci_out_step = exci_out_step[:,4*200:-(rpad*200),:]
    #    hifi_out_list.append(hifi_out_step)
    #    exci_out_list.append(exci_out_step)

    #hifi0_out = tf.concat(hifi_out_list,axis=1)
    #exci_out = tf.concat(exci_out_list,axis=1)
    #print('hifi0 {} {}'.format(hifi0_out.shape,exci_out.shape))

    hifi1_win = 21
    hifi1_olp = 3
    hifi1_up = 10
    hifi_out_list = []
    time_len = tf.shape(hifi0_out).numpy()[1]
    for step in range(hifi1_olp,time_len,hifi1_win-2*hifi1_olp):
        start = step-hifi1_olp
        end = step+hifi1_win-hifi1_olp
        melstep = hifi0_out[:,start:end,:]
        excistep = exci_out[:,start*200:end*200]
        rpad = hifi1_win - tf.shape(melstep).numpy()[1]
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
            excistep = tf.concat([excistep,tf.zeros([1,rpad*200,excistep.shape[2]])],axis=1)
        hifi_out_step = hifi1_model((melstep,excistep),training=True)
        print('hifi_out {} '.format(hifi_out_step.shape))
        if step == hifi1_olp:
            hifi_out_step = hifi_out_step[:,:-30,:]
        else:
            if end<=time_len:
                hifi_out_step = hifi_out_step[:,30:-30,:]
            else:
                hifi_out_step = hifi_out_step[:,30:-(rpad*10):,:]
        hifi_out_list.append(hifi_out_step)

    hifi1_out = tf.concat(hifi_out_list,axis=1)
    print('hifi1 {}'.format(hifi1_out.shape))

    hifi2_win = 150
    hifi2_olp = 5
    hifi2_up = 5
    hifi_out_list = []
    time_len = tf.shape(hifi1_out).numpy()[1]
    for step in range(hifi2_olp,time_len,hifi2_win-2*hifi2_olp):
        start = step-hifi2_olp
        end = step+hifi2_win-hifi2_olp
        melstep = hifi1_out[:,start:end,:]
        excistep = exci_out[:,(start)*20:(end)*20]
        rpad = hifi2_win - tf.shape(melstep).numpy()[1]
        print('rpad {} {} {} {}'.format(step,rpad,melstep.shape,excistep.shape))
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
            excistep = tf.concat([excistep,tf.zeros([1,rpad*20,excistep.shape[2]])],axis=1)
        hifi_out_step = hifi2_model((melstep,excistep),training=True)
        #print('hifi_out {} '.format(hifi_out_step.shape))
        if step == hifi2_olp:
            hifi_out_step = hifi_out_step[:,:-25,:]
        else:
            if end<=time_len:
                hifi_out_step = hifi_out_step[:,25:-25,:]
            else:
                if rpad > 0:
                    hifi_out_step = hifi_out_step[:,25:-(rpad*5):,:]
                else:
                    hifi_out_step = hifi_out_step[:,25:,:]

        print('hifi_out {} '.format(hifi_out_step.shape))
        hifi_out_list.append(hifi_out_step)

    hifi2_out = tf.concat(hifi_out_list,axis=1)
    print('hifi2 {}'.format(hifi2_out.shape))

    hifi3_win = 700
    hifi3_olp = 20
    hifi3_up = 2
    hifi_out_list = []
    time_len = tf.shape(hifi2_out).numpy()[1]
    for step in range(hifi3_olp,time_len,hifi3_win-2*hifi3_olp):
        start = step-hifi3_olp
        end = step+hifi3_win-hifi3_olp
        melstep = hifi2_out[:,start:end,:]
        excistep = exci_out[:,start*4:end*4]
        rpad = hifi3_win - tf.shape(melstep).numpy()[1]
        print('rpad {} {} {}'.format(rpad,melstep.shape,excistep.shape))
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
            excistep = tf.concat([excistep,tf.zeros([1,rpad*4,excistep.shape[2]])],axis=1)
        hifi_out_step = hifi3_model((melstep,excistep),training=True)
        print('hifi_out {} '.format(hifi_out_step.shape))
        if step == hifi3_olp:
            hifi_out_step = hifi_out_step[:,:-40,:]
        else:
            if end<=time_len:
                hifi_out_step = hifi_out_step[:,40:-40,:]
            else:
                hifi_out_step = hifi_out_step[:,40:-(rpad*2):,:]
        hifi_out_list.append(hifi_out_step)

    hifi3_out = tf.concat(hifi_out_list,axis=1)
    print('hifi3 {}'.format(hifi3_out.shape))

    hifi4_win = 1320 
    hifi4_olp = 20
    hifi4_up = 2
    hifi_out_list = []
    time_len = tf.shape(hifi3_out).numpy()[1]
    for step in range(hifi4_olp,time_len,hifi4_win-2*hifi4_olp):
        start = step-hifi4_olp
        end = step+hifi4_win-hifi4_olp
        melstep = hifi3_out[:,start:end,:]
        excistep = exci_out[:,(start)*2:(end)*2]
        rpad = hifi4_win - tf.shape(melstep).numpy()[1]
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
            excistep = tf.concat([excistep,tf.zeros([1,rpad*2,excistep.shape[2]])],axis=1)
        hifi_out_step = hifi4_model((melstep,excistep),training=True)
        print('hifi_out {} '.format(hifi_out_step.shape))
        if step == hifi4_olp:
            hifi_out_step = hifi_out_step[:,:-40,:]
        else:
            if end <= time_len:
                hifi_out_step = hifi_out_step[:,40:-40,:]
            else:
                hifi_out_step = hifi_out_step[:,40:-(rpad*2):,:]
        hifi_out_list.append(hifi_out_step)

    hifi4_out = tf.concat(hifi_out_list,axis=1)
    print('hifi4 {}'.format(hifi4_out.shape))

    hifi_out_list = []
    time_len = tf.shape(hifi4_out).numpy()[1]
    for step in range(4,time_len,8000):
        melstep = hifi4_out[:,step-4:step+8004,:]
        rpad = 8008 - tf.shape(melstep).numpy()[1]
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
        hifi_out_step = hifi5_model((melstep),training=True)
        print('hifi_out {} '.format(hifi_out_step.shape))
        if step == 4:
            hifi_out_step = hifi_out_step[:,:-4,:]
        else:
            if rpad == 0:
                hifi_out_step = hifi_out_step[:,4:-4,:]
            else:
                hifi_out_step = hifi_out_step[:,4:-(rpad):,:]
        hifi_out_list.append(hifi_out_step)

    hifi5_out = tf.concat(hifi_out_list,axis=1)
    print('hifi5 {}'.format(hifi5_out.shape))
    pcm = (tf.reshape(hifi5_out,(-1))*32768.0).numpy().astype(dtype=np.int16)
    import wave
    with wave.open("tfdebug.wav", 'wb') as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(16000)
        wf.writeframes(pcm.data)

    #hifi0_model.save('tfmodel/hifigan0')
    #hifi1_model.save('tfmodel/hifigan1')
    #hifi2_model.save('tfmodel/hifigan2')
    #hifi3_model.save('tfmodel/hifigan3')
    #hifi4_model.save('tfmodel/hifigan4')
    #hifi5_model.save('tfmodel/hifigan5')



#voc_infer(model_id.voices['F7'].tmp_vocinput)
